#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_mjn_sample_qstat.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/18 13:59
@Description :
    生成欧盟报告所需的样本质控信息（sample.qstat）。
    主要统计测序数据量、人类序列数、微生物比对情况等指标。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import pandas as pd


# ----------------------------------------------------------------------
# 参数解析
# ----------------------------------------------------------------------
def _parse_args() -> Namespace:
    """
    解析命令行参数。

    Returns
    -------
    Namespace
        包含 job_id、sample_id 和 out_file 的参数对象。
    """
    parser = ArgumentParser(description="生成欧盟报告的样本质控统计信息")

    parser.add_argument(
        "--job_id",
        type=str,
        required=True,
        help="芯片编号，用于定位分析目录"
    )
    parser.add_argument(
        "--sample_id",
        type=str,
        required=True,
        help="样本编号，用于查找样本的质控文件"
    )
    parser.add_argument(
        "--out_file",
        type=str,
        required=True,
        help="输出文件路径（TSV 格式）"
    )

    return parser.parse_args()


# ----------------------------------------------------------------------
# SAM 计数函数
# ----------------------------------------------------------------------
def _sam_reads(sam_tmp: Path) -> int:
    """
    统计 SAM 文件中的唯一序列数（read name 去重）。

    Parameters
    ----------
    sam_tmp : Path
        SAM 文件路径。

    Returns
    -------
    int
        序列数量。若文件不存在或为空，则返回 0。
    """
    if not sam_tmp.is_file() or sam_tmp.stat().st_size == 0:
        return 0

    reads = set()
    with open(sam_tmp, "r") as sam_file:
        for line in sam_file:
            line = line.strip()
            if not line or line.startswith("@"):
                continue
            reads.add(line.split()[0])

    return len(reads)


# ----------------------------------------------------------------------
# 主逻辑
# ----------------------------------------------------------------------
def main(args: Namespace) -> None:
    """
    读取 sample.qstat 文件并统计各类比对结果，生成用于欧盟报告的统计表。

    Parameters
    ----------
    args : Namespace
        解析命令行参数后的对象。
    """
    job_id = args.job_id.strip()
    sample_id = args.sample_id.strip()

    # 创建输出目录
    out_file = Path(args.out_file).absolute()
    out_file.parent.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # 加载 sample.qstat 文件
    # ------------------------------------------------------------------
    qstat_file = Path("/data/mNGS/runmngs/tmp/report") / f"{sample_id}.sample.qstat.xls"
    if not qstat_file.exists():
        raise FileNotFoundError(f"找不到 sample.qstat 文件：{qstat_file}")

    df = pd.read_csv(qstat_file, sep="\t", thousands=",")

    # 筛选对应 job_id + sample_id 的行
    df = df.loc[(df["jobid"] == job_id) & (df["sample_id"] == sample_id)]
    if df.empty:
        raise ValueError(f"在 {qstat_file} 中找不到匹配 job_id={job_id}, sample_id={sample_id} 的记录")

    row = df.iloc[0]  # 单样本必定只有一行

    # ------------------------------------------------------------------
    # 统计比对信息
    # ------------------------------------------------------------------
    sample_dir = Path("/data/mNGS/runmngs/result") / job_id / sample_id

    out_list = [
        ["欧蒙样本编号", sample_id],
        ["测序所得序列数(raw reads)", row["allreads"]],
        ["测序所得序列数(clean reads)", "NA"],
        ["测序所得序列数(uniq reads)", "NA"],
        ["Q30（%）", f"{row['q30']}%"],
        ["比对到人基因组上的序列数", row["humanreads"]]
    ]
    # 全部微生物序列
    all_microbial_reads = _sam_reads(sample_dir / f"{sample_id}_merged.sam")
    out_list.append(["比对到微生物基因组的序列数", all_microbial_reads])

    # 细菌序列（6 个文件合并）
    bac_files = [sample_dir / f"prokaryotes_{i}.sam_tmp" for i in range(6)]
    bac_reads = sum(_sam_reads(f) for f in bac_files)
    out_list.append(["比对到细菌基因组的序列数", bac_reads])

    # 真菌序列
    fungi_reads = _sam_reads(sample_dir / "Fungi.sam_tmp")
    out_list.append(["比对到真菌基因组的序列数", fungi_reads])

    # 病毒序列
    virus_reads = _sam_reads(sample_dir / "virus.sam_tmp")
    out_list.append(["比对到病毒基因组的序列数", virus_reads])

    # 寄生虫序列
    parasite_reads = _sam_reads(sample_dir / "Parasite.sam_tmp")
    out_list.append(["比对到寄生虫基因组的序列数", parasite_reads])

    # ------------------------------------------------------------------
    # 写出文件
    # ------------------------------------------------------------------
    with open(out_file, "w") as handle:
        for item in out_list:
            handle.write("\t".join(map(str, item)) + "\n")


# ----------------------------------------------------------------------
if __name__ == "__main__":
    main(_parse_args())
